import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

# 检测GPU是否可用
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"使用设备: {device}")

# 定义一个简单的神经网络模型
class SimpleNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        """
        初始化神经网络
        input_size: 输入特征的维度
        hidden_size: 隐藏层神经元数量
        output_size: 输出层神经元数量
        """
        super(SimpleNN, self).__init__()
        self.layer1 = nn.Linear(input_size, hidden_size)  # 第一层：输入层到隐藏层
        self.relu = nn.ReLU()  # 激活函数
        self.layer2 = nn.Linear(hidden_size, output_size)  # 第二层：隐藏层到输出层
        self.sigmoid = nn.Sigmoid()  # 添加sigmoid激活函数用于二分类

    def forward(self, x):
        """
        前向传播
        x: 输入数据
        """
        x = self.layer1(x)
        x = self.relu(x)
        x = self.layer2(x)
        x = self.sigmoid(x)  # 使用sigmoid进行二分类
        return x

def calculate_accuracy(outputs, targets):
    """
    计算准确率
    outputs: 模型输出
    targets: 真实标签
    """
    predictions = (outputs >= 0.5).float()
    correct = (predictions == targets).sum().item()
    total = targets.size(0)
    return correct / total

# 生成示例训练数据
def generate_sample_data(num_samples=1000):
    """
    生成示例数据用于训练
    num_samples: 样本数量
    """
    # 生成随机输入数据 (这里假设是二维特征)
    X = np.random.randn(num_samples, 2)
    # 生成目标值 (这里使用简单的规则：如果x1 + x2 > 0，则为1，否则为0)
    y = np.array([1 if x1 + x2 > 0 else 0 for x1, x2 in X])
    return torch.FloatTensor(X).to(device), torch.FloatTensor(y).view(-1, 1).to(device)

def train_model(model, X_train, y_train, epochs=100, learning_rate=0.01):
    """
    训练模型
    model: 神经网络模型
    X_train: 训练数据
    y_train: 训练标签
    epochs: 训练轮数
    learning_rate: 学习率
    """
    criterion = nn.BCELoss()  # 使用二元交叉熵损失函数
    optimizer = optim.SGD(model.parameters(), lr=learning_rate)  # 使用随机梯度下降优化器

    for epoch in range(epochs):
        # 前向传播
        outputs = model(X_train)
        loss = criterion(outputs, y_train)

        # 计算准确率
        accuracy = calculate_accuracy(outputs, y_train)

        # 反向传播和优化
        optimizer.zero_grad()  # 清除之前的梯度
        loss.backward()  # 计算梯度
        optimizer.step()  # 更新参数

        # 每10轮打印一次损失值和准确率
        if (epoch + 1) % 10 == 0:
            print(f'Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}, Accuracy: {accuracy:.4f}')

# 示例使用
if __name__ == "__main__":
    # 创建模型实例
    input_size = 2
    hidden_size = 10
    output_size = 1
    model = SimpleNN(input_size, hidden_size, output_size).to(device)

    # 生成训练数据
    X_train, y_train = generate_sample_data()

    # 训练模型
    train_model(model, X_train, y_train)

    # 测试模型
    with torch.no_grad():
        test_input = torch.FloatTensor([[1.0, 2.0]]).to(device)  # 测试样例
        prediction = model(test_input)
        print(f"测试输入 [1.0, 2.0] 的预测结果: {prediction.item():.4f}")
        print(f"预测类别: {'正类' if prediction.item() >= 0.5 else '负类'}")
